"""
YASA分析性能测试

本文件用于测试YASA分析器处理20分钟数据的性能表现。
测试流程：
1. 生成20分钟的模拟EEG数据（614,400个数据点，512Hz采样率）
2. 将数据分割成40个epochs（每个30秒）
3. 直接调用YASA分析器进行睡眠分期分析
4. 测量YASA分析的总耗时

原理：
直接使用项目中现有的YASAAnalyzer类，生成足够的数据让其进行完整的睡眠分析。
数据生成时间不计入分析时间，只测量纯YASA分析耗时。
"""

import sys
import os
import time
import numpy as np
from datetime import datetime, timedelta
from typing import List, Tuple

# 添加项目根目录到Python路径
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..'))
sys.path.insert(0, project_root)

from app.core.yasa_analyzer import YASAAnalyzer


class YASAPerformanceTester:
    """YASA性能测试器"""
    
    def __init__(self):
        self.sampling_rate = 512  # 512Hz采样率
        self.duration_minutes = 20  # 20分钟
        self.total_samples = self.duration_minutes * 60 * self.sampling_rate  # 614,400个数据点
        self.epoch_duration = 30  # 30秒一个epoch
        self.total_epochs = (self.duration_minutes * 60) // self.epoch_duration  # 40个epochs
        
        print(f"YASA性能测试初始化:")
        print(f"- 采样率: {self.sampling_rate}Hz")
        print(f"- 测试时长: {self.duration_minutes}分钟")
        print(f"- 总数据点: {self.total_samples:,}个")
        print(f"- 总epochs: {self.total_epochs}个")
        print("-" * 50)
    
    def generate_mock_eeg_data(self) -> Tuple[np.ndarray, List[datetime]]:
        """
        生成20分钟的模拟EEG数据
        
        Returns:
            tuple: (数据数组, 时间戳列表)
        """
        print("开始生成20分钟模拟EEG数据...")
        start_time = time.time()
        
        # 生成基础时间戳
        base_time = datetime.now()
        timestamps = []
        for i in range(self.total_samples):
            timestamp = base_time + timedelta(seconds=i / self.sampling_rate)
            timestamps.append(timestamp)
        
        # 生成模拟EEG信号
        # 使用多种频率成分模拟真实的脑电信号
        t = np.linspace(0, self.duration_minutes * 60, self.total_samples)
        
        # 基础信号：多种频率成分的组合
        delta_wave = 10 * np.sin(2 * np.pi * 2 * t)      # Delta波 (1-4Hz)
        theta_wave = 8 * np.sin(2 * np.pi * 6 * t)       # Theta波 (4-8Hz)
        alpha_wave = 6 * np.sin(2 * np.pi * 10 * t)      # Alpha波 (8-13Hz)
        beta_wave = 4 * np.sin(2 * np.pi * 20 * t)       # Beta波 (13-30Hz)
        
        # 添加随机噪声模拟真实环境
        noise = np.random.normal(0, 2, self.total_samples)
        
        # 组合信号
        eeg_data = delta_wave + theta_wave + alpha_wave + beta_wave + noise
        
        # 添加一些睡眠特征（模拟不同睡眠阶段的特征）
        for i in range(0, len(eeg_data), self.sampling_rate * 60):  # 每分钟变化
            end_idx = min(i + self.sampling_rate * 60, len(eeg_data))
            # 模拟不同睡眠阶段的振幅变化
            stage_factor = np.random.uniform(0.5, 1.5)
            eeg_data[i:end_idx] *= stage_factor
        
        generation_time = time.time() - start_time
        print(f"数据生成完成，耗时: {generation_time:.2f}秒")
        print(f"数据形状: {eeg_data.shape}")
        print(f"时间戳数量: {len(timestamps)}")
        
        return eeg_data, timestamps
    
    def run_yasa_analysis_test(self) -> dict:
        """
        运行YASA分析性能测试
        
        Returns:
            dict: 测试结果
        """
        print("\n" + "="*60)
        print("开始YASA分析性能测试")
        print("="*60)
        
        # 1. 生成数据（不计入分析时间）
        eeg_data, timestamps = self.generate_mock_eeg_data()
        
        # 2. 初始化YASA分析器
        print("\n初始化YASA分析器...")
        analyzer = YASAAnalyzer(
            sampling_rate=self.sampling_rate,
            channel_names=['EEG']
        )
        
        # 3. 开始YASA分析计时
        print("\n开始YASA分析...")
        analysis_start_time = time.time()
        
        # 4. 批量添加数据进行分析
        # 将数据分批添加，模拟实际的数据流
        batch_size = self.sampling_rate * 30  # 每次添加30秒的数据
        analysis_results = []
        
        for i in range(0, len(eeg_data), batch_size):
            end_idx = min(i + batch_size, len(eeg_data))
            batch_data = eeg_data[i:end_idx]
            batch_timestamps = timestamps[i:end_idx]
            
            print(f"处理数据批次: {i//batch_size + 1}/{self.total_epochs}")
            
            # 添加数据批次
            results = analyzer.add_data_batch(batch_data, batch_timestamps)
            if results:
                analysis_results.extend(results)
        
        # 5. 结束计时
        analysis_end_time = time.time()
        total_analysis_time = analysis_end_time - analysis_start_time
        
        # 6. 获取最终分析结果
        final_result = analyzer.get_current_sleep_stage()
        analysis_history = analyzer.get_analysis_history_for_ui()
        
        # 7. 整理测试结果
        test_result = {
            'test_duration_minutes': self.duration_minutes,
            'total_data_points': self.total_samples,
            'total_epochs': self.total_epochs,
            'analysis_time_seconds': total_analysis_time,
            'analysis_time_minutes': total_analysis_time / 60,
            'data_points_per_second': self.total_samples / total_analysis_time,
            'epochs_per_second': self.total_epochs / total_analysis_time,
            'final_sleep_stage': final_result,
            'total_analysis_results': len(analysis_results),
            'analysis_history_count': len(analysis_history),
            'yasa_available': analyzer.is_yasa_available()
        }
        
        return test_result
    
    def print_test_results(self, results: dict):
        """
        打印测试结果
        
        Args:
            results: 测试结果字典
        """
        print("\n" + "="*60)
        print("YASA分析性能测试结果")
        print("="*60)
        
        print(f"测试数据:")
        print(f"  - 测试时长: {results['test_duration_minutes']}分钟")
        print(f"  - 总数据点: {results['total_data_points']:,}个")
        print(f"  - 总epochs: {results['total_epochs']}个")
        
        print(f"\n性能指标:")
        print(f"  - YASA分析总耗时: {results['analysis_time_seconds']:.2f}秒")
        print(f"  - YASA分析总耗时: {results['analysis_time_minutes']:.2f}分钟")
        print(f"  - 数据处理速度: {results['data_points_per_second']:.0f}个数据点/秒")
        print(f"  - Epoch处理速度: {results['epochs_per_second']:.2f}个epochs/秒")
        
        print(f"\n分析结果:")
        print(f"  - YASA库可用: {'是' if results['yasa_available'] else '否'}")
        print(f"  - 最终睡眠阶段: {results['final_sleep_stage']}")
        print(f"  - 分析结果数量: {results['total_analysis_results']}")
        print(f"  - 历史记录数量: {results['analysis_history_count']}")
        
        # 计算效率指标
        if results['analysis_time_minutes'] > 0:
            efficiency_ratio = results['test_duration_minutes'] / results['analysis_time_minutes']
            print(f"\n效率指标:")
            print(f"  - 实时处理倍数: {efficiency_ratio:.1f}x")
            if efficiency_ratio >= 1.0:
                print(f"  - ✅ 可以实时处理 (处理速度 >= 1x)")
            else:
                print(f"  - ❌ 无法实时处理 (处理速度 < 1x)")


def main():
    """主函数"""
    try:
        # 创建测试器
        tester = YASAPerformanceTester()
        
        # 运行测试
        results = tester.run_yasa_analysis_test()
        
        # 打印结果
        tester.print_test_results(results)
        
        print("\n" + "="*60)
        print("测试完成！")
        print("="*60)
        
    except Exception as e:
        print(f"\n❌ 测试过程中发生错误: {str(e)}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()